#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2022/7/13 下午5:14
# @Author  : qiufengfeng
# @File    : cls_mlp_model.py
# @Description :
from typing import Dict,Any


from tensorflow.keras.layers import Dropout,Lambda,Dense
from tensorflow import keras
from nlp_tools.tasks.classification.abc_model import ABCClassificationModel
class ClsMlpModel(ABCClassificationModel):
    @classmethod
    def default_hyper_parameters(cls) -> Dict[str, Dict[str, Any]]:
        return {
            'layer_output': {

            }
        }

    def build_model_arc(self) -> None:
        output_dim = self.label_processor.vocab_size

        config = self.hyper_parameters
        embed_model = self.embedding.embed_model

        # build model structure in sequent way
        layer_stack = [
            Lambda(lambda x: x[:, 0]),
            Dropout(rate=0.3),
            Dense(output_dim),
            self._activation_layer()
        ]

        tensor = embed_model.output
        for layer in layer_stack:
            tensor = layer(tensor)

        self.tf_model: keras.Model = keras.Model(embed_model.inputs, tensor)
